import torch
import torch_geometric.transforms as T
from torch.nn import functional as F
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import HGTConv, Linear


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
        super().__init__()

        self.lin_dict = torch.nn.ModuleDict()
        for node_type in data.node_types:
            self.lin_dict[node_type] = Linear(-1, hidden_channels)

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HGTConv(hidden_channels, hidden_channels, data.metadata(),
                           num_heads, group='sum')
            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)


    def forward(self, x_dict, edge_index_dict):         # x_dict:  {node_type: [num_nodes: node_feature]}, edge_index_dict: {edge_type: [2: num_edges]}
        for node_type, x in x_dict.items():                            # 每种
            x_dict[node_type] = self.lin_dict[node_type](x).relu_()         # 先将node_feature 降维到 hidden_channels

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)          # HGTConv 的返回  out_dict[node_type] = out        x_dict[node_type] : (num_node, hidden_channels)

        return self.lin(x_dict['author'])           # 节点预测 我们这里是图的预测，

model = HGT(hidden_channels=64, out_channels=dataset.num_classes,
            num_heads=2, num_layers=2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)


def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x_dict, data.edge_index_dict)
    mask = data['paper'].train_mask
    loss = F.cross_entropy(out['paper'][mask], data['paper'].y[mask])
    loss.backward()
    optimizer.step()
    return float(loss)

train()